#!/bin/bash

# ==============================================================================
# 自动检测并导出现有的Docker镜像脚本
# 根据docker images的输出智能检测相关镜像并导出
# ==============================================================================

# 创建输出目录
OUTPUT_DIR="docker_images_backup"
mkdir -p $OUTPUT_DIR

# 定义输出文件名（包含时间戳）
OUTPUT_FILE="$OUTPUT_DIR/yuxi_know_existing_images_$(date +%Y%m%d_%H%M%S).tar"

echo "==============================================="
echo "自动检测并导出现有的 Docker 镜像"
echo "==============================================="
echo "输出文件: $OUTPUT_FILE"
echo ""

# 定义要搜索的镜像关键字
declare -A IMAGE_PATTERNS=(
    ["yuxi"]="yuxi-api yuxi-web"
    ["neo4j"]="neo4j"
    ["postgres"]="postgres postgres-for-rag"
    ["etcd"]="etcd"
    ["milvus"]="milvus"
    ["minio"]="minio"
    ["python"]="python"
    ["uv"]="ghcr.io/astral-sh/uv"
    ["node"]="node"
    ["nginx"]="nginx"
)

echo "正在扫描本地镜像..."
echo "-----------------------------------------------"

# 获取所有本地镜像并提取相关镜像
declare -A FOUND_IMAGES
ALL_LOCAL_IMAGES=$(docker images --format "{{.Repository}}:{{.Tag}}" | grep -v "<none>")

for IMAGE in $ALL_LOCAL_IMAGES; do
    # 检查是否匹配我们的模式
    for PATTERN in "${!IMAGE_PATTERNS[@]}"; do
        if echo "$IMAGE" | grep -qi "$PATTERN"; then
            # 检查是否在期望的镜像列表中
            EXPECTED_IMAGES="${IMAGE_PATTERNS[$PATTERN]}"
            for EXPECTED in $EXPECTED_IMAGES; do
                if echo "$IMAGE" | grep -q "$EXPECTED"; then
                    FOUND_IMAGES["$IMAGE"]=1
                    echo "✓ 找到相关镜像: $IMAGE"
                    break 2
                fi
            done
        fi
    done
done

# 特殊处理一些已知存在的镜像
SPECIAL_IMAGES=(
    "quay.io/coreos/etcd:v3.5.5"
    "milvusdb/milvus:v2.5.6"
    "minio/minio:RELEASE.2023-03-20T20-16-18Z"
)

for IMAGE in "${SPECIAL_IMAGES[@]}"; do
    if echo "$ALL_LOCAL_IMAGES" | grep -q "^$IMAGE$"; then
        FOUND_IMAGES["$IMAGE"]=1
        echo "✓ 找到特殊镜像: $IMAGE"
    fi
done

if [ ${#FOUND_IMAGES[@]} -eq 0 ]; then
    echo "❌ 没有找到相关的镜像"
    echo ""
    echo "可用的本地镜像:"
    echo "$ALL_LOCAL_IMAGES" | head -20
    echo ""
    echo "请确保Docker容器已构建，或者手动指定要导出的镜像"
    exit 1
fi

# 准备导出列表
EXPORT_IMAGES=()
for IMAGE in "${!FOUND_IMAGES[@]}"; do
    EXPORT_IMAGES+=("$IMAGE")
done

echo ""
echo "找到 ${#EXPORT_IMAGES[@]} 个相关镜像："
echo "-----------------------------------------------"
for i in "${!EXPORT_IMAGES[@]}"; do
    echo "$((i+1)). ${EXPORT_IMAGES[$i]}"
done
echo ""

# 显示镜像大小信息
echo "镜像大小信息："
echo "-----------------------------------------------"
for IMAGE in "${EXPORT_IMAGES[@]}"; do
    SIZE=$(docker images --format "{{.Repository}}:{{.Tag}}\t{{.Size}}" | grep "^$IMAGE" | cut -f2)
    echo "$IMAGE - $SIZE"
done
echo ""

# 确认是否继续
read -p "确认导出这些镜像吗？(y/N): " -n 1 -r
echo ""
if [[ ! $REPLY =~ ^[Yy]$ ]]; then
    echo "导出已取消"
    exit 0
fi

# 保存镜像到 tar 文件
echo "开始保存镜像到 tar 文件..."
echo "-----------------------------------------------"

echo "正在保存 ${#EXPORT_IMAGES[@]} 个镜像到 $OUTPUT_FILE..."
if docker save ${EXPORT_IMAGES[@]} -o "$OUTPUT_FILE"; then
    echo "✓ 镜像保存成功"
else
    echo "❌ 镜像保存失败"
    exit 1
fi

# 计算文件大小
FILE_SIZE=$(du -h "$OUTPUT_FILE" | cut -f1)

# 生成镜像列表文件
LIST_FILE="$OUTPUT_DIR/existing_image_list_$(date +%Y%m%d_%H%M%S).txt"
echo "生成镜像列表文件: $LIST_FILE"
echo "Yuxi-Know 现有Docker镜像导出列表" > "$LIST_FILE"
echo "导出时间: $(date)" >> "$LIST_FILE"
echo "镜像文件: $OUTPUT_FILE" >> "$LIST_FILE"
echo "文件大小: $FILE_SIZE" >> "$LIST_FILE"
echo "" >> "$LIST_FILE"
echo "包含的镜像:" >> "$LIST_FILE"
for IMAGE in "${EXPORT_IMAGES[@]}"; do
    SIZE=$(docker images --format "{{.Repository}}:{{.Tag}}\t{{.Size}}" | grep "^$IMAGE" | cut -f2)
    echo "  - $IMAGE ($SIZE)" >> "$LIST_FILE"
done

# 生成加载脚本
LOAD_SCRIPT="$OUTPUT_DIR/load_existing_images.sh"
echo "生成加载脚本: $LOAD_SCRIPT"
cat > "$LOAD_SCRIPT" << 'EOF'
#!/bin/bash

# Docker 镜像加载脚本
# 使用方法: ./load_existing_images.sh [image_file.tar]

# 默认镜像文件
IMAGE_FILE=""
if [ $# -eq 1 ]; then
    IMAGE_FILE="$1"
else
    # 查找最新的 tar 文件
    IMAGE_FILE=$(ls -t docker_images_backup/*existing*.tar 2>/dev/null | head -1)
fi

if [ -z "$IMAGE_FILE" ] || [ ! -f "$IMAGE_FILE" ]; then
    echo "错误: 找不到镜像文件"
    echo "使用方法: $0 <image_file.tar>"
    exit 1
fi

echo "正在从 $IMAGE_FILE 加载 Docker 镜像..."
echo ""

if docker load -i "$IMAGE_FILE"; then
    echo ""
    echo "✓ 镜像加载成功!"
    echo ""
    echo "加载的镜像列表:"
    docker images --format "table {{.Repository}}:{{.Tag}}\t{{.Size}}\t{{.CreatedAt}}" | head -n 20
else
    echo "❌ 镜像加载失败"
    exit 1
fi
EOF

chmod +x "$LOAD_SCRIPT"

# 最终输出
echo ""
echo "==============================================="
echo "导出完成！"
echo "==============================================="
echo "✓ 镜像文件: $OUTPUT_FILE"
echo "✓ 文件大小: $FILE_SIZE"
echo "✓ 镜像数量: ${#EXPORT_IMAGES[@]}"
echo "✓ 镜像列表: $LIST_FILE"
echo "✓ 加载脚本: $LOAD_SCRIPT"
echo ""
echo "使用方法:"
echo "  1. 加载镜像: $LOAD_SCRIPT"
echo "  2. 或者直接: docker load -i $OUTPUT_FILE"
echo "==============================================="